import cv2
from PIL import Image
import numpy as np
import random
import math


class Transform:
    """分类Transform的基类
    """

    def __init__(self):
        pass


class Compose(Transform):
    """根据数据预处理/增强算子对输入数据进行操作。
       所有操作的输入图像流形状均是[H, W, C]，其中H为图像高，W为图像宽，C为图像通道数。
    Args:
        transforms (list): 数据预处理/增强算子。
    Raises:
        TypeError: 形参数据类型不满足需求。
        ValueError: 数据长度不匹配。
    """

    def __init__(self, transforms):
        if not isinstance(transforms, list):
            raise TypeError('The transforms must be a list!')
        if len(transforms) < 1:
            raise ValueError('The length of transforms ' + \
                             'must be equal or larger than 1!')
        self.transforms = transforms

        # 检查transforms里面的操作，目前支持PaddleX定义的或者是imgaug操作
        for op in self.transforms:
            if not isinstance(op, Transform):
                import imgaug.augmenters as iaa
                if not isinstance(op, iaa.Augmenter):
                    raise Exception(
                        "Elements in transforms should be defined in 'paddlex.cls.transforms' or class of imgaug.augmenters.Augmenter, see docs here: https://paddlex.readthedocs.io/zh_CN/latest/apis/transforms/"
                    )

    def __call__(self, im, label=None):
        """
        Args:
            im (str/np.ndarray): 图像路径/图像np.ndarray数据。
            label (int): 每张图像所对应的类别序号。
        Returns:
            tuple: 根据网络所需字段所组成的tuple；
                字段由transforms中的最后一个数据预处理操作决定。
        """
        if isinstance(im, np.ndarray):
            if len(im.shape) != 3:
                raise Exception(
                    "im should be 3-dimension, but now is {}-dimensions".
                        format(len(im.shape)))
        else:
            try:
                # im = cv2.imread(im)

                im = Image.open(im).convert('RGB')
                im = np.array(im)
            except:
                raise TypeError('Can\'t read The image file {}!'.format(im))
        # im = im.astype('float32')
        # im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
        for op in self.transforms:
            if isinstance(op, Transform):
                outputs = op(im, label)
                im = outputs[0]
                if len(outputs) == 2:
                    label = outputs[1]
            else:
                import imgaug.augmenters as iaa
                if isinstance(op, iaa.Augmenter):
                    im = execute_imgaug(op, im)
                outputs = (im,)
                if label is not None:
                    outputs = (im, label)
        return outputs

    def add_augmenters(self, augmenters):
        if not isinstance(augmenters, list):
            raise Exception(
                "augmenters should be list type in func add_augmenters()")
        transform_names = [type(x).__name__ for x in self.transforms]
        for aug in augmenters:
            if type(aug).__name__ in transform_names:
                logging.error(
                    "{} is already in ComposedTransforms, need to remove it from add_augmenters().".
                        format(type(aug).__name__))
        self.transforms = augmenters + self.transforms


class Normalize(Transform):
    """对图像进行标准化。

    1. 对图像进行归一化到区间[0.0, 1.0]。
    2. 对图像进行减均值除以标准差操作。

    Args:
        mean (list): 图像数据集的均值。默认为[0.485, 0.456, 0.406]。
        std (list): 图像数据集的标准差。默认为[0.229, 0.224, 0.225]。

    """

    def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
        self.mean = mean
        self.std = std

    def __call__(self, im, label=None):
        """
        Args:
            im (np.ndarray): 图像np.ndarray数据。
            label (int): 每张图像所对应的类别序号。

        Returns:
            tuple: 当label为空时，返回的tuple为(im, )，对应图像np.ndarray数据；
                   当label不为空时，返回的tuple为(im, label)，分别对应图像np.ndarray数据、图像类别id。
        """
        mean = np.array(self.mean)
        std = np.array(self.std)
        # im = normalize(im, mean, std)
        im = im / 255.
        im -= mean
        im /= std
        if label is None:
            return (im, )
        else:
            return (im, label)


class RandomHorizontalFlip(Transform):
    """以一定的概率对图像进行随机水平翻转，模型训练时的数据增强操作。

    Args:
        prob (float): 随机水平翻转的概率。默认为0.5。
    """

    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, im, label=None):
        """
        Args:
            im (np.ndarray): 图像np.ndarray数据。
            label (int): 每张图像所对应的类别序号。

        Returns:
            tuple: 当label为空时，返回的tuple为(im, )，对应图像np.ndarray数据；
                   当label不为空时，返回的tuple为(im, label)，分别对应图像np.ndarray数据、图像类别id。
        """
        if random.random() < self.prob:
            im = self.horizontal_flip(im)

        if label is None:
            return (im,)
        else:
            return (im, label)

    def horizontal_flip(self, im):
        if len(im.shape) == 3:
            im = im[:, ::-1, :]
        elif len(im.shape) == 2:
            im = im[:, ::-1]
        return im


class RandomCropProb(Transform):
    """对图像进行随机剪裁，模型训练时的数据增强操作。

    1. 根据lower_scale、lower_ratio、upper_ratio计算随机剪裁的高、宽。
    2. 根据随机剪裁的高、宽随机选取剪裁的起始点。
    3. 剪裁图像。
    4. 调整剪裁后的图像的大小到crop_size*crop_size。

    Args:
        prob (float): 使用随机裁剪的概率，小于这个概率会使用裁剪，否则返回原图，默认为0.5。
        crop_size (int): 随机裁剪后重新调整的目标边长。默认为224。
        lower_scale (float): 裁剪面积相对原面积比例的最小限制。默认为0.08。
        lower_ratio (float): 宽变换比例的最小限制。默认为3. / 4。
        upper_ratio (float): 宽变换比例的最大限制。默认为4. / 3。
    """

    def __init__(self,
                 prob=0.5,
                 crop_size=224,
                 lower_scale=0.08,
                 lower_ratio=3. / 4,
                 upper_ratio=4. / 3):
        self.crop_size = crop_size
        self.lower_scale = lower_scale
        self.lower_ratio = lower_ratio
        self.upper_ratio = upper_ratio
        self.prob = prob

    def __call__(self, im, label=None):
        """
        Args:
            im (np.ndarray): 图像np.ndarray数据。
            label (int): 每张图像所对应的类别序号。

        Returns:
            tuple: 当label为空时，返回的tuple为(im, )，对应图像np.ndarray数据；
                   当label不为空时，返回的tuple为(im, label)，分别对应图像np.ndarray数据、图像类别id。
        """
        if random.random() < self.prob:
            im = self.random_crop(im, self.crop_size, self.lower_scale,
                                  self.lower_ratio, self.upper_ratio)
        if label is None:
            return (im,)
        else:
            return (im, label)

    def random_crop(self, im,
                    crop_size=224,
                    lower_scale=0.08,
                    lower_ratio=3. / 4,
                    upper_ratio=4. / 3):
        scale = [lower_scale, 1.0]
        ratio = [lower_ratio, upper_ratio]
        aspect_ratio = math.sqrt(np.random.uniform(*ratio))
        w = 1. * aspect_ratio
        h = 1. / aspect_ratio
        bound = min((float(im.shape[0]) / im.shape[1]) / (h ** 2),
                    (float(im.shape[1]) / im.shape[0]) / (w ** 2))
        scale_max = min(scale[1], bound)
        scale_min = min(scale[0], bound)
        target_area = im.shape[0] * im.shape[1] * np.random.uniform(
            scale_min, scale_max)
        target_size = math.sqrt(target_area)
        w = int(target_size * w)
        h = int(target_size * h)
        i = np.random.randint(0, im.shape[0] - h + 1)
        j = np.random.randint(0, im.shape[1] - w + 1)
        im = im[i:i + h, j:j + w, :]
        im = cv2.resize(im, (crop_size, crop_size))
        return im


class Resize(Transform):
    def __init__(self, size):
        self.size = size

    def __call__(self, im, label=None):
        im = cv2.resize(im, (self.size, self.size))
        if label is None:
            return (im,)
        else:
            return (im, label)
